import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision.models import resnet50

# 投影头（用于对比学习）
class ProjectionHead(nn.Module):
    def __init__(self, in_dim=2048, hidden_dim=2048, out_dim=128):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(in_dim, hidden_dim),
            nn.BatchNorm1d(hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, out_dim)
        )

    def forward(self, x):
        return self.layers(x)

# SimCLR模型（融合监督信号）
class SimCLR(nn.Module):
    def __init__(self, backbone='resnet50', projection_dim=128):
        super().__init__()
        # 骨干网络（ResNet50）
        self.backbone = resnet50(pretrained=False)
        self.backbone.fc = nn.Identity()  # 移除原始全连接层
        self.feature_dim = 2048  # ResNet50输出维度
        
        # 投影头
        self.projection_head = ProjectionHead(
            in_dim=self.feature_dim,
            out_dim=projection_dim
        )

    def forward(self, x):
        # x: [batch_size * num_views, C, H, W]
        features = self.backbone(x)  # [batch_size * num_views, 2048]
        projections = self.projection_head(features)  # [batch_size * num_views, 128]
        return features, projections

# 监督+无监督融合的InfoNCE损失
def simclr_loss(projections, labels, temperature=0.07):
    """
    projections: [N, D] 其中N = batch_size * num_views
    labels: [batch_size] 原始样本标签（每个样本有num_views个投影）
    """
    batch_size = labels.shape[0]
    num_views = projections.shape[0] // batch_size  # 每个样本的视图数
    N = projections.shape[0]  # 添加这行，定义N变量
    
    # 计算余弦相似度矩阵
    sim = F.cosine_similarity(
        projections.unsqueeze(1), 
        projections.unsqueeze(0), 
        dim=2
    )  # [N, N]
    
    # 构建掩码：同一原始样本的不同视图为正例
    mask = torch.eye(batch_size, device=projections.device).repeat(num_views, num_views)
    mask = mask - torch.eye(N, device=projections.device)  # 排除自身
    
    # 监督掩码：同一类别的样本也视为正例
    label_mask = (labels.unsqueeze(0) == labels.unsqueeze(1)).float()  # [B, B]
    label_mask = label_mask.repeat(num_views, num_views)  # [N, N]
    mask = mask + label_mask  # 融合两种正例掩码
    mask = mask.clamp(0, 1)  # 确保掩码非负
    
    # 计算InfoNCE损失
    sim = sim / temperature
    exp_sim = torch.exp(sim)
    exp_sim = exp_sim * (1 - torch.eye(N, device=projections.device))  # 排除自身
    
    log_prob = sim - torch.log(exp_sim.sum(dim=1, keepdim=True))
    loss = (-mask * log_prob).sum(dim=1) / mask.sum(dim=1)  # 平均正例损失
    return loss.mean()